import argparse
import warnings
import joblib
import numpy as np
import pandas as pd
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, LabelEncoder, OneHotEncoder
from sklearn.metrics import roc_auc_score, classification_report, confusion_matrix
from imblearn.over_sampling import SMOTE
import optuna
from optuna.samplers import TPESampler
from lightgbm import LGBMClassifier
from lightgbm.callback import early_stopping
from xgboost import XGBClassifier
import matplotlib.pyplot as plt

plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['font.size'] = 15


class PowerLoadPredict(object):
    def __init__(self, file_path):
        self.data_source = pd.read_csv(file_path)
        self.model = None
        self.feature_columns = None

    def load_model(self, model_path, feature_columns_path=None):
        """加载模型和特征列"""
        self.model = joblib.load(model_path)
        if feature_columns_path:
            self.feature_columns = joblib.load(feature_columns_path)
        print('模型加载成功')

    def feature_engineering(self, data):
        """
        特征工程 - 需要与训练集保持一致
        """
        feature_data = data.copy()

        # 使用与训练集相同的特征
        one_data = feature_data[["OverTime", "MaritalStatus", "DistanceFromHome", "JobRole",
                                "Department", "PerformanceRating", "Gender", "JobInvolvement",
                                'JobSatisfaction', 'StockOptionLevel', 'YearsWithCurrManager',
                                'YearsInCurrentRole', 'MonthlyIncome', 'YearsAtCompany',
                                'JobLevel', 'Age', 'TotalWorkingYears']]

        # 进行one-hot编码
        feature_data_processed = pd.get_dummies(one_data)

        # 如果已经加载了训练时的特征列，确保测试集特征与训练集一致
        if self.feature_columns is not None:
            # 添加训练集中有但测试集中没有的特征列，并填充0
            for col in self.feature_columns:
                if col not in feature_data_processed.columns and col != 'Attrition':
                    feature_data_processed[col] = 0

            # 只保留训练集中存在的特征列，并按相同顺序排列
            feature_data_processed = feature_data_processed[self.feature_columns[1:]]  # 排除Attrition列

        return feature_data_processed

    def predict(self, data=None):
        """进行预测"""
        if data is None:
            data = self.data_source

        # 特征工程
        X_test = self.feature_engineering(data)

        # 预测
        y_pred_proba = self.model.predict_proba(X_test)[:, 1]  # 预测概率
        y_pred = self.model.predict(X_test)  # 预测类别

        return y_pred, y_pred_proba

    def evaluate_model(self, y_true, y_pred, y_pred_proba):
        """评估模型性能"""
        if y_true is not None:
            # 计算AUC
            auc_score = roc_auc_score(y_true, y_pred_proba)
            print(f"测试集AUC: {round(auc_score * 100, 2)}%")

            # 分类报告
            print("\n分类报告:")
            print(classification_report(y_true, y_pred))

            # 混淆矩阵
            cm = confusion_matrix(y_true, y_pred)
            print("混淆矩阵:")
            print(cm)

            return auc_score
        else:
            print("未提供真实标签，无法计算评估指标")
            return None

    def save_predictions(self, predictions, output_path='../results/predictions.csv'):
        """保存预测结果"""
        results_df = pd.DataFrame({
            '预测类别': predictions[0],
            '预测概率': predictions[1]
        })

        # 确保输出目录存在
        import os
        os.makedirs(os.path.dirname(output_path), exist_ok=True)

        results_df.to_csv(output_path, index=False, encoding='utf-8-sig')
        print(f'预测结果已保存到: {output_path}')

        return results_df


def main():
    # 初始化预测类
    plp = PowerLoadPredict('../data/test.csv')

    # 加载模型和特征列
    plp.load_model('../model/model.pkl')

    # 如果你在训练时保存了特征列，可以这样加载
    # plp.load_model('../model/model.pkl', '../model/feature_columns.pkl')

    # 进行预测
    predictions = plp.predict()

    # 保存预测结果
    results_df = plp.save_predictions(predictions)

    # 如果你有测试集的真实标签，可以进行评估
    # 假设测试集中有'Attrition'列
    if 'Attrition' in plp.data_source.columns:
        # 对真实标签进行编码（如果需要）
        le = LabelEncoder()
        y_true = le.fit_transform(plp.data_source['Attrition'])
        plp.evaluate_model(y_true, predictions[0], predictions[1])

    # 显示前几条预测结果
    print("\n前10条预测结果:")
    print(results_df.head(10))


if __name__ == '__main__':
    main()